#!python

# -*- coding: utf-8 -*-

import sys, os
import socket
import popen2
import thread
import select
import thread
import threading
import time
import traceback
import base64
import hashlib
import struct

from collections import deque

g_debug = True

g_exit = False
clients = []
address = {}
message = {}
filelist = {}
alias = {}
talkto = {}
talktoalias = {}
typemask = {}
def send_msg(s, msg_bytes):
    if typemask[s] == "WSS":
        #switch gbk back to unicode and then to utf-8
        msg_bytes = unicode(msg_bytes,"gbk").encode("utf-8")
        token = b"\x81"
        length = len(msg_bytes)
        if length < 126:
            token += struct.pack("B", length)
        elif length <= 0xFFFF:
            token += struct.pack("!BH", 126, length)
        else:
            token += struct.pack("!BQ", 127, length)
        msg = token + msg_bytes
        s.send(msg)
    else:
        s.send(msg_bytes)
def handledata(s):
    try:
        data = s.recv(1024)
        if data[data.find("Upgrade:")+9:][:9] == "websocket":
            typemask[s] = "WSS"
            #for websocket handshake
            header_dict = {}

            header, body = data.encode("utf-8").split("\r\n\r\n")
            header_list = header.split("\r\n")
            for i in range(0, len(header_list)):
                if i == 0:
                    if len(header_list[i].split(" ")) == 3:
                        header_dict["method"], header_dict["url"], header_dict["protocol"] = header_list[i].split(" ")
                else:
                    k, v = header_list[i].split(":", 1)
                    header_dict[k]= v.strip()
            response_tpl = "HTTP/1.1 101 Switching Protocols\r\n"\
                            "Upgrade:websocket\r\n"\
                            "Sec-WebSocket-Accept:%s\r\n"\
                            "Connection:Upgrade\r\n"\
                            "Sec-WebSocket-Origin:null\r\n"\
                            "Sec-WebSocket-Location:ws://%s%s\r\n\r\n"
            magic_key = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
            value = header_dict["Sec-WebSocket-Key"] + magic_key;
            ac = base64.b64encode(hashlib.sha1(value.encode("utf-8")).digest())
            response_str = response_tpl % (ac.decode("utf-8"), header_dict["Host"], header_dict["url"])
            #print(bytes(response_str.encode("utf-8")))
            print(str(address[s]) + "WebSocket connection")
            s.send(bytes(response_str.encode("utf-8")))
        else:
            if typemask[s] == "WSS":
                try:
                    #unpack wss message

                    if len(data) > 1:
                        payload_len = ord(data[1])&127
                        if payload_len == 126:
                            extend_payload_len = data[2:4]
                            mask = data[4:8]
                            decoded = data[8:]
                        elif payload_len == 127:
                            extend_payload_len = data[2:10]
                            mask = data[10:14]
                            decoded = data[14:]
                        else:
                            extend_payload_len = None
                            mask = data[2:6]
                            decoded = data[6:]

                        bytes_list = bytearray()
                        for i in range(len(decoded)):
                            chunk = ord(decoded[i]) ^ ord(mask[i % 4])
                            bytes_list.append(chunk)
                        body = bytes_list.decode("utf-8")

                        if body == "close":
                            #disconnect
                            disconnect(s)
                        else:
                            #inner code is gbk
                            message[s].append(body.encode("gbk"))
                        
                except:
                    if g_debug:
                        traceback.print_exc()
                        print "Error on :" + str(address[s])
                    #disconnect
                    disconnect(s)
            else:
                if data == '':
                    if typemask[s] == "WSS":
                        pass
                    else:
                        #disconnect
                        disconnect(s)
                else:
                    message[s].append(data)
    except:
        #if g_debug:
            traceback.print_exc()
            print "Error on :" + str(address[s])
def disconnect(s):
    print(str(address[s]) + "client disconnected")
    clients.remove(s)
    #inform other clients
    for _s in clients:
        send_msg(_s, "disconnected " + alias[s] + " " + str(address[s]))
    for _s in talkto:
        if talkto[_s] == s:
            talkto[_s] = _s;
            send_msg(_s, "talkto "  + alias[s] + " " + str(address[s]) + " failed")
    del talktoalias[s]
    del message[s]
    del filelist[s]
    del address[s]
    del alias[s]
    del talkto[s]
    del typemask[s]
    s.close()
    
def datahandler(s):
    while True:
        try:
            if g_exit:
                break

            rlist, wlist, xlist = select.select([s,],[],[], 0)
            if rlist:
                clientsocket, clientaddress = s.accept()
                #inform other clients
                for _s in clients:
                    send_msg(_s, "connected " + str(clientaddress))
                clients.append(clientsocket)
                message[clientsocket] = deque()
                address[clientsocket] = clientaddress
                alias[clientsocket] = str(clientaddress)
                talkto[clientsocket] = clientsocket
                talktoalias[clientsocket] = None
                typemask[clientsocket] = None
                filelist[clientsocket]= deque()
                print(str(clientaddress) + "client connected!")
            if len(clients) > 0:
                rlist, wlist, xlist = select.select(clients, [], [], 0)
                if rlist:
                    for i in rlist:
                        handledata(i)
                        
            time.sleep(0.1)
        except:
            if g_debug:
                traceback.print_exc()
    print("Thread Exited") 
def msghandler():
    while True:
        try:
            if g_exit:
                break
            for s in message:
                if len(message[s]) > 0:
                    data = message[s].pop()
                    if data:
                        print(str(address[s]) + data)
                        param = data.split(" ")
                        #do something
                        if param[0] == "users":
                            result = ""
                            for a in address:
                                result += str(address[a]) + "[" + alias[a] + "]" + "\n"
                            print("user list send to " +str(address[s]))
                            send_msg(s, result)
                        elif param[0] == "me":
                            result = str(address[s]) + "[" + alias[s] + "]" + "\n"
                            send_msg(s, result)
                        elif param[0] == "alias":
                            res = ""
                            if len(param) > 1:
                                talk = None
                                for key in alias:
                                    if alias[key] == param[1]:
                                        talk = key
                                        break
                                if talk is not None:
                                    res = str(address[s]) + "duplicated alias!"
                                else:
                                    alias[s] = param[1]
                                    res = str(address[s]) + "aliased to " + param[1] + " success."
                            else:
                                res = str(address[s]) + "alias failed!"
                            send_msg(s, res)
                            print(res)
                        elif param[0] == "talk":
                            res = ""
                            talk = None
                            if len(param) > 1:
                                for key in alias:
                                    if alias[key] == param[1]:
                                        talk = key
                                        break
                            if talk is not None:
                                res = alias[s] + " is talking to " + alias[key] + "."
                            else:
                                res = alias[s] + " talk faild!"
                            talkto[s] = talk
                            talktoalias[s] = alias[key]
                            send_msg(s, res)
                            print(res)
                        elif param[0] == "send":
                            talk = None
                            if len(param) > 2:
                                for key in alias:
                                    if alias[key] == param[1]:
                                        talk = key
                                        break;
                            if talk is not None:
                                send_msg(talk, alias[s] + ":" + " ".join(param[2:]))
                        elif param[0] == "all":
                            if len(param) > 1:
                                for _s in clients:
                                    send_msg(_s, alias[s] + ":" + " ".join(param[1:]))
                        elif param[0] == "file":
                            if len(param) > 1:
                                talk = talkto[s]
                                #restore talkto
                                if talk is None or talk == s:
                                   if talktoalias[s] is not None:
                                       for key in alias:
                                           if alias[key] == talktoalias[s]:
                                               talk = key
                                               break
                                   talkto[s] = talk
                                    
                                status = int(param[1])
                                filesize = None
                                filepath = None
                                index = 0
                                
                                if status == 1:
                                    if len(param) > 3:
                                        name = param[3]
                                        if name:      
                                            for key in alias:
                                                if alias[key] == name:
                                                    talk = key
                                                    break;
                                        talkto[s] = talk
                                if status == 0:
                                    if len(param) > 4:
                                        name = param[4]
                                        if name:      
                                            for key in alias:
                                                if alias[key] == name:
                                                    talk = key
                                                    break;
                                        talkto[s] = talk
                                        
                                if talk is None or talk == s:
                                    print("No talking subject")
                                else:
                                    print("File header")

                                    if status == 1:
                                        if len(param) > 2:
                                            index = int(param[2])
                                            filesize = filelist[s][index]["size"]
                                            filepath = filelist[s][index]["path"]                                
                                    elif status == 3:
                                        if len(param) > 2:
                                            index = int(param[2])
                                            filesize = filelist[s][index]["size"]
                                            filepath = filelist[s][index]["path"]
                                    if status == 0:
                                        if len(param) > 3:
                                            filesize = param[2]
                                            filepath = param[3]
                                            filehead = {}
                                            filehead["path"] = filepath
                                            filehead["size"] = filesize
                                            filelist[talk].append(filehead)
                                            index = len(filelist[talk]) - 1
                                    elif status == 2:
                                        if len(param) > 2:
                                            index = int(param[2])
                                            filesize = filelist[talk][index]["size"]
                                            filepath = filelist[talk][index]["path"]
                                    send_msg(talk, alias[s] + ":file " + str(status) + " " + str(index) + " " + filesize + " " + filepath)
                                        
                        #elif param[0] == "file":
                        #    if len(param) > 1:
                        #        talk = talkto[s]
                        #        #restore talkto
                        #        if talk is None or talk == s:
                        #           if talktoalias[s] is not None:
                        #               for key in alias:
                        #                   if alias[key] == talktoalias[s]:
                        #                       talk = key
                        #                       break
                        #           talkto[s] = talk
                        #        if talk is None or talk == s:
                        #            print("No talking subject")
                        #        else:
                        #            print("File header")
                        #            send_msg(talk, "file:" + alias[s] + " " + " ".join(param[1:]))
                        elif param[0] == "data":
                            if len(param) > 1:
                                talk = talkto[s]
                                #restore talkto
                                if talk is None or talk == s:
                                   if talktoalias[s] is not None:
                                       for key in alias:
                                           if alias[key] == talktoalias[s]:
                                               talk = key
                                               break
                                   talkto[s] = talk
                                if talk is None or talk == s:
                                    print("No talking subject")
                                else:
                                    print("File data")
                                    send_msg(talk, data)
                        else:
                             talk = talkto[s]
                             #restore talkto
                             if talk is None or talk == s:
                                if talktoalias[s] is not None:
                                    for key in alias:
                                        if alias[key] == talktoalias[s]:
                                            talk = key
                                            break
                                talkto[s] = talk
                             if talk is not None:
                                 send_msg(talk, alias[s] + ":" + data)
                        

            time.sleep(0.5)
        except:
            if g_debug:
                traceback.print_exc()
    print("msg Thread Exited")  

def main():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind((socket.gethostname(), 8000))
    
    s.listen(10)
    
    dt = threading.Thread(target=datahandler, args=(s,))
    dt.start()
    mt = threading.Thread(target=msghandler, args=())
    mt.start()
    
    try:
        while True:
            content = raw_input("")
            if content == "exit":
                break
            #do something else
            elif content == "users":
                result = ""
                for a in address:
                    result += str(address[a]) + "[" + alias[a] + "]" + "\n"
                print(result)
            else:
                param = content.split(" ")
                if param[0] == "all":
                    if len(param) > 1:
                        for s in clients:
                            send_msg(s, "server:" + " ".join(param[1:]))
                    
            

            time.sleep(0.5)
    except:
        if g_debug:
            traceback.print_exc()
    global g_exit
    g_exit = True
    dt.join()
    mt.join()
    s.close()
    print("Exited")
                
if __name__ == "__main__":
    main()
